Make Emily.Fast fused kernels composable inside a caller's defn (fixes #205) - #206
Open
lostbean wants to merge 3 commits into
Open
Make Emily.Fast fused kernels composable inside a caller's defn (fixes #205)#206lostbean wants to merge 3 commits into
lostbean wants to merge 3 commits into
Conversation
Every defn-callable Emily.Fast kernel (rms_norm, layer_norm, rope, rope_with_freqs, and both scaled_dot_product_attention variants) was a plain def, so calling one from inside a defn-defined function raised "cannot invoke ... inside defn because it was not defined with defn" — defn rewrites remote calls to Nx.Defn.Compiler.__remote__/4, which dispatches to the callee's generated __defn:name__/arity, and only defn/deftransform definitions export that. The kernels only worked as the literal top-level unit handed to Nx.Defn.jit_apply, contradicting the moduledoc's own composition example. Convert the public kernels to deftransform: the body stays plain Elixir (Nx.block/4's pin-matched callback isn't expressible inside a defn body, and the *_fallback helpers are applied as ordinary Elixir on Expr parameters by Nx.Defn.Expr.block/4 at trace time), while the generated __defn:name__ delegator makes the call dispatchable from any caller's defn. The eager path and the fused mx::fast::* dispatch are unchanged. Adds composed_defn_test covering the two-level nesting repro from the issue, per-kernel fused-vs-fallback conformance of the composed form, a positive assertion that composed graphs lower to the fast_* IR opcodes, and a negative assertion that no [:emily, :block, :fallback] telemetry fires for Emily.Fast blocks on the evaluator lane. Fixes ausimian#205. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sdpa_fallback/sdpa_sinks_fallback built their causal bias from
less_equal(iota({q_len, 1}), iota({1, k_len})) — i.e. query i allowed
to attend keys j >= i, the *anti*-causal direction — while
mx::fast::scaled_dot_product_attention's mask_mode "causal" masks the
standard direction with bottom-right alignment (query i attends keys
j <= i + (k_len - q_len), see make_or_fetch_mask in mlx/fast.cpp). The
fused path and the composed fallback therefore disagreed on every
causal: true call.
This went unnoticed because the existing sdpa causal conformance test
runs with Nx.default_backend(Emily.Backend) in setup, which routes its
"Evaluator oracle" lane through the fused kernel as well — the
comparison was fused-vs-fused. The composed conformance tests added
for ausimian#205 run the oracle lane on a genuine Nx.BinaryBackend process and
caught the divergence.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Emily.Fast.rope_with_freqs/4 documents (and Emily.Bumblebee.FastKernels supplies) the HF inverse-frequency convention: theta = position * freqs. mx::fast::rope expects the reciprocal table — it computes theta = position / freqs (reciprocal(inputs[2]) in mlx/fast.cpp's fallback, and 1.0 / freqs[...] in the Metal kernels). Both native lanes passed the table through untouched, so every position > 0 rotated by the wrong angles on the fused path while the composed fallback rotated correctly — the two paths only agreed at position 0, which is exactly what the existing offset-0 conformance test exercised (and its oracle lane ran fused anyway, see the previous commit). Take the reciprocal before handing the table to the NIF, in both the eager Emily.Backend.block dispatch and the Expr-compiler IR lowering. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #205.
What
Three commits, smallest-possible diffs:
36fbb04— makeEmily.Fastfused kernels callable inside a caller's owndefn. Converts the eight public kernels (rms_norm/3,layer_norm/4,rope/3,rope_with_freqs/4, and the fourscaled_dot_product_attention*variants) from plaindeftodeftransform, plus a newtest/emily/fast/composed_defn_test.exs(14 tests).einsumstaysdef(eager-only); every*_fallbackstays plaindefp.e32e3c4— correct the causal-mask direction in the SDPA composed fallback. The mask was inverted (attended future keys). Now bottom-right-aligned to matchmx::fast::sdpa: query rowiattends keysj <= i + (k_len - q_len).5ade402— invert the freqs table beforemx::fast::ropein both native lanes. emily documents the HF convention (θ = position × inv_freq) but MLX's rope computes θ = position / freqs, so the fusedrope_with_freqswas wrong at every position > 0. BothEmily.Backend.block/4and theEmily.CompilerIR lowering now pass the reciprocal.Why
deftransform, not the issue's suggesteddefp→defnpThe failure is not in the fallbacks.
defnrewrites remote calls to dispatch through a generated__defn:name__/aritydelegator that onlydefn/deftransformdefinitions export — a plain-defkernel raises "was not defined with defn" before its body runs, which is why only the literal top-leveljit_apply(fn ...)form ever worked.deftransformexports the delegator while keeping the body plain Elixir, which is required:Nx.block/4's callback is applied as ordinary Elixir at trace time (Nx.Defn.Expr.block/4) and pin-matches the block struct (fn ^block, ... ->), which is illegal inside defn. TheNx.blockconstruction and the entire fast path are byte-for-byte untouched.Bugs 2 and 3 were exposed by the new conformance tests and are pre-existing (present in 1.0.0), not composition regressions — the composed and top-level lanes agreed with each other on the wrong values. They went unnoticed because the existing per-kernel tests set
Nx.default_backend(Emily.Backend)insetup, so their oracle lane also runs fused and the comparison is fused-vs-fused (diff exactly 0.0). I left those tests as-is to keep this PR scoped, but rebasing their oracles ontoNx.BinaryBackendis worth a follow-up — happy to split any of the three commits into separate PRs if preferred.Test coverage (
composed_defn_test.exs)defn→defnp→Emily.Fast.*) forrms_norm,rope, andscaled_dot_product_attention_with_mask, including under globalEmily.Backend+Emily.Compilerdefaults.compiler: Emily.Compiler, native: true, native_fallback: :raisevsNx.Defn.EvaluatoronNx.BinaryBackend, within f32 tolerance.:fast_*IR opcodes, the fused lane runs withnative_fallback: :raise, and a telemetry handler flunks on any[:emily, :block, :fallback]event for anEmily.Fast.Block.*struct.Validation
lerobot/smolvla_basecheckpoint) from per-layer eager calls to one composeddefnkept numerical parity (0.646% MRE vs the Python reference, unchanged) and cut warm inference from ~1.26s to ~611ms median — the composed-graph speedup this issue is after.Note: all changes are Elixir-only; the native suite was run against the released, checksum-verified 1.0.0 NIF (my dev machine lacks the Metal toolchain for an MLX source build). Nothing under
c_src/is touched.🤖 Generated with Claude Code